#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import pysolr
from pymongo import MongoClient
import re
import sys
import time


def PickTermOrNot(line):
    line = line.strip()
    pattern = re.compile("^(.{7,}|[\w\-]{4,6}|(?!\d+$)(?![a-zA-Z]+$)[\da-zA-Z]{3})$")
    m = re.match(pattern, line)
    if m:
        return True
    else:
        return False


if __name__ == "__main__":
    time1 = time.time()
    solr = pysolr.Solr('http://localhost:8080/solr/plosp_pdf')

    SERVER = 'localhost'
    PORT = 27017
    DB_NAME = "KnowledgeBase"
    COLLECTION_NAME = "KnowledgeBase_Test"

    connection = MongoClient(SERVER, PORT)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    infile = open(sys.argv[1], "r")
    headline = infile.readline()
    splitheadline = headline.split("\t")
    seq = [0, 1, 2, 8, 9, 10, 11]
    for line in infile:
        line = line.rstrip("\n")
        if line.find("withdrawn") != -1:
            continue
        splitline = line.split("\t")
        query = 'fileContent:('
        for i in seq:
            if splitline[i] != "":
                if (splitline[i].startswith("\"") and splitline[i].endswith("\"")):
                    result1s = splitline[i][1:-1].split("|")
                    for result1 in result1s:
                        if (result1 != ""):
                            if PickTermOrNot(result1):
                                query += "\"" + result1 + "\" OR "
                else:
                    if PickTermOrNot(splitline[i]):
                        query += "\"" + splitline[i] + "\" OR "
        query = query[:-4] + ')'
        results1 = solr.search(query)
        results2 = solr.search(query, start=0, rows=results1.hits)
        count = 0
        for result in results2:
            if (result["score"] >= 3):
                collection.update({"PMID": result["pmid"]}, {"$addToSet": {"HGNC": splitline[0]}}, True)
                count += 1
            else:
                break
                # print splitline[0] + "\t" + str(results1.hits) + "\t" + str(count)
                # print query
                # print "############################################################"

    infile.close()
    connection.close()
    time2 = time.time()
    print str(time2 - time1)
